# mat文件处理

from pathlib import Path
import numpy as np
import scipy.io as scio
import torch
import torch.nn.functional as F


def read_sensor_data(file_path: Path, transform=None, modality='US'):
    rf_data = scio.loadmat(str(file_path))
    sensor_data = rf_data['rf_data']
    # if modality == 'US':
    #     sensor_data = rf_data['rf_data']
    #     #acoustic_impedance = rf_data['acoustic_impedance']
    # elif modality == 'PA':
    #     sensor_data = rf_data['rf_data']
    #     #optical_absorption = rf_data['optical_absorption']
    # else:
    #     raise NotImplementedError("input modality is not supported!")

    if transform is not None:
        sensor_data = transform(sensor_data)
    return sensor_data

def preprocess_sensor_data(sensor_data, target_shape=(2560, 64), resize_method='pad', pad_value=-1.0):
    """
    传感器数据预处理函数，将数据归一化到 [-1, 1] 范围
    
    参数:
    - sensor_data: 原始传感器数据 numpy 数组
    - target_shape: 目标形状 (宽, 高)
    - resize_method: 调整尺寸的方法，'interpolate' 或 'pad'
    
    返回:
    - 预处理后的张量，形状为 [1, target_shape[0], target_shape[1]]，范围为 [-1, 1]
    """
    # 转换为张量
    tensor_data = torch.from_numpy(sensor_data).float() if isinstance(sensor_data, np.ndarray) else torch.tensor(sensor_data, dtype=torch.float32)
    
    # 处理维度 - 确保是 3D 张量 [C, H, W]
    if len(tensor_data.shape) == 2:
        # 如果是 [H, W] 格式，添加通道维度变成 [1, H, W]
        tensor_data = tensor_data.unsqueeze(0)
    elif len(tensor_data.shape) == 3 and tensor_data.shape[0] != 1:
        # 如果是 [H, W, C] 格式，转换为 [C, H, W]
        tensor_data = tensor_data.permute(2, 0, 1)
    
    # 获取当前形状
    current_shape = tensor_data.shape[1:]  # [H, W]
    
    # 归一化到 [-1, 1] 范围
    min_val = tensor_data.min()
    max_val = tensor_data.max()
    
    # 避免除零错误
    if max_val > min_val:
        # 先归一化到 [0, 1]，再缩放到 [-1, 1]
        tensor_data = (tensor_data - min_val) / (max_val - min_val)
        tensor_data = tensor_data * 2 - 1
    else:
        # 如果数据是常数，设置为0
        tensor_data = torch.zeros_like(tensor_data)
        print(f"Data is constant. Setting to 0 !!!!!")
    
    # 如果尺寸不匹配，根据选择的方法进行调整
    if current_shape != target_shape:
        if resize_method == 'interpolate':
            tensor_data = F.interpolate(
                tensor_data.unsqueeze(0),  # 添加批次维度
                size=target_shape,
                mode='bilinear',
                align_corners=False
            ).squeeze(0)  # 移除批次维度
        elif resize_method == 'pad':
            # 使用填充方式调整尺寸
            target_h, target_w = target_shape
            current_h, current_w = current_shape
            
            # 计算填充量，将数据居中放置
            pad_h = max(0, target_h - current_h)
            pad_w = max(0, target_w - current_w)
            
            pad_top = pad_h // 2
            pad_bottom = pad_h - pad_top
            pad_left = pad_w // 2
            pad_right = pad_w - pad_left
            
            # 如果原始数据比目标尺寸大，需要裁剪
            if current_h > target_h or current_w > target_w:
                crop_h_start = max(0, (current_h - target_h) // 2)
                crop_w_start = max(0, (current_w - target_w) // 2)
                crop_h_end = crop_h_start + min(target_h, current_h)
                crop_w_end = crop_w_start + min(target_w, current_w)
                
                tensor_data = tensor_data[:, crop_h_start:crop_h_end, crop_w_start:crop_w_end]
                
                # 重新计算填充量
                current_h, current_w = tensor_data.shape[1:]
                pad_h = max(0, target_h - current_h)
                pad_w = max(0, target_w - current_w)
                pad_top = pad_h // 2
                pad_bottom = pad_h - pad_top
                pad_left = pad_w // 2
                pad_right = pad_w - pad_left
            
            # 应用填充
            tensor_data = F.pad(tensor_data, (pad_left, pad_right, pad_top, pad_bottom), value=pad_value)
        else:
            raise ValueError(f"Unsupported resize_method: {resize_method}. Use 'interpolate' or 'pad'.")
    
    # 确保数据类型正确
    return tensor_data.float()
